# %%
import pandas as pd
import numpy as np
from openai import OpenAI
from tqdm import tqdm

import start

prompts_file = start.MAIN_DIR + "prompts_characterization.xlsx"
prompts = pd.read_excel(prompts_file, sheet_name=None)


text_file = start.CLEAN_DIR + "character_classifications_gold.xlsx"
text_df = pd.read_excel(text_file)
text_df = text_df[text_df.text.notnull()]
text_df = text_df[text_df.set == "test"]

client = OpenAI(api_key=start.OPENAI_API_KEY)
MODEL = "gpt-4o"

# %%
# Specify the desired prompt name, or leave empty for all
SHEET_NAME = "Few Shot 3a"

sheet_data = prompts[SHEET_NAME]

prompt = []
for message in sheet_data.index:
    dict_entry = {
        "role": sheet_data.loc[message, "role"],
        "content": sheet_data.loc[message, "content"],
    }
    prompt.append(dict_entry)

responses = []
for text in tqdm(text_df.text):
    messages = prompt + [{"role": "user", "content": text}]

    response = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        temperature=0.00,
    )

    cleaned_response = response.choices[0].message.content
    responses.append(cleaned_response)

text_df["response"] = responses
text_df["agreement"] = np.where(text_df.response == text_df.character_gold, 1, 0)
text_df = text_df[
    [
        "unique_id",
        "text",
        "set",
        "response",
        "character_gold",
        "agreement",
    ]
]
text_df.to_excel(
    start.MAIN_DIR + f"data/clean/gpt_classifications_characters_test.xlsx",
    index=False,
)
